# -*- coding: utf-8 -*-

# File Name： vis
# Description :
# Author : lirui
# create_date： 2022/6/6
# Change Activity:
from enum import Enum

import cv2
import numpy as np


class Color(Enum):
    """An enum that defines common colors.

    Contains red, green, blue, cyan, yellow, magenta, white and black.
    """
    red = (0, 0, 255)
    green = (0, 255, 0)
    blue = (255, 0, 0)
    cyan = (255, 255, 0)
    yellow = (0, 255, 255)
    magenta = (255, 0, 255)
    white = (255, 255, 255)
    black = (0, 0, 0)
    gray = (122, 114, 114)


def draw_bbox(img, bboxes, color=Color.cyan.value, thickness=2):
    """
    draw bounding box to image. bbox with shape (width_min,height_min,w_max,h_max)
    Args:
        img:
        bboxes: every box with shape (width_min,height_min,w_max,h_max)
        color:
        thickness:

    Returns:

    """
    tmp_img = img.copy()
    for bbox in bboxes:
        bbox = np.array(bbox)
        bbox = bbox.reshape(-1, 2)
        ws, hs = bbox[:, 0], bbox[:, 1]
        x0, x1, y0, y1 = ws.min(), ws.max(), hs.min(), hs.max()
        points = [[x0, y0], [x1, y0], [x1, y1], [x0, y1]]
        result = np.array(points).astype(np.int32)
        vector = result.reshape(-1, 2)
        cv2.polylines(tmp_img, [vector], isClosed=True, color=color, thickness=thickness)
    return tmp_img


def paint_contours(img, contours, draw_value=255, paint_channel=1):
    """

    Args:
        paint_channel:
        img:
        contours:
        draw_value:

    Returns:

    """
    tmp_img = img.copy()
    tmp_img = tmp_img.astype(dtype=np.int32)
    bg = np.zeros(tmp_img.shape[:2], dtype=np.int32)
    bg = cv2.drawContours(bg, contours, cv2.FILLED, draw_value, -1)
    a = tmp_img[:, :, paint_channel] + bg
    a[np.where(a > 255)] = 255
    tmp_img[:, :, paint_channel] = a
    return tmp_img


def draw_contours(img, contours, color=Color.cyan.value, thickness=2):
    """

    Args:
        img:
        contours: list or ndarray, the first channel must be batch size. (batch_size,points_size,x_y)
        color:
        thickness:

    Returns:

    """
    tmp_img = img.copy()
    cv2.polylines(tmp_img, contours, isClosed=True, color=color, thickness=thickness)
    return tmp_img
